from copy import deepcopy
from auto_optimizer import OnnxGraph
from .pass_manager import PASS_MANAGER
from .onnx_passes import Status


class OnnxPassExecuter:
    def __init__(self, pass_list=[]):
        self.pass_list = []
        self.add_passes(pass_list)

    def add_passes(self, passes, allow_repeat=False):
        for pass_func in passes:
            if pass_func in self.pass_list and not allow_repeat:
                continue
            self.pass_list.append(pass_func)

    def del_passes(self, passes):
        for pass_func in passes:
            if pass_func in self.pass_list:
                self.pass_list.remove(pass_func)
            else:
                print(f'Not contain the pass function, skip it: {pass_func}')

    def execute(self, onnx_path, save_path=None):
        if not save_path:
            save_path = onnx_path
        graph = OnnxGraph.parse(onnx_path)

        is_changed = False
        for pass_func in self.pass_list:
            status_code, graph = pass_func(graph)
            if status_code == Status.SUCCESS:
                is_changed = True
            elif status_code == Status.NOCHANGE:
                is_changed = True if is_changed else False
            else:
                raise Exception(f'Excute onnx pass failed: {pass_func}')

        if is_changed:
            graph.save(save_path)
            return Status.SUCCESS
        else:
            return Status.NOCHANGE

    @classmethod
    def execute_passes(cls, onnx_path, pass_list=[], save_path=None):
        executer = cls(pass_list=pass_list)
        return executer.execute(onnx_path, save_path=save_path)


def build_pass_executer(model_type):
    pass_list = PASS_MANAGER.get_pass(model_type)
    pass_executer = OnnxPassExecuter(pass_list=pass_list)
    return pass_executer
